'''
Created on Jul 3, 2009

@author: Roni
'''
import cliques
import cliques.mdp.heuristicMdp
import math
class mdpLookahead(cliques.mdp.heuristicMdp.heuristicMdp):
    '''
    classdocs
    '''
    def __init__(self,k,edge_prob, num_of_nodes):
        '''
        Constructor
        '''
        cliques.mdp.heuristicMdp.heuristicMdp.__init__(self, k, edge_prob, num_of_nodes)    
        self.lookahead_depth=2


    def add_state(self, state, v, openlist, moves, action, next_state, next_state_prob):
        cliques.mdp.heuristicMdp.heuristicMdp.add_state(self, state, v, openlist, moves, action, next_state, next_state_prob)
        
        if len(self.states) > self.upper_bound_states:
            alert("OVER UPPER BOUND - BUG!!")

    def set_lookahead_depth(self,new_depth):
        self.lookahead_depth=new_depth
        
        # Calculate upper bound of #states
        self.upper_bound_states = 1
        nodes = int(self.nodes)
        for i in xrange(nodes-1,max(nodes-1-new_depth,0),-1):
            self.upper_bound_states=self.upper_bound_states*i*pow(2,i)
        print "O(#states)=%d" % self.upper_bound_states
    
    
    
    def expand(self, state, v, policy, openlist):
        moves = state.g
        if self.is_goal(state):
            v[state] = moves
            policy[state] = state
            self.goals.add(state)
        else:
            actions = state.actions()
            if len(actions)==0: # No relevant actions
                v[state]=moves+self.h(state)
                policy[state]=None
            else: # Has actions - need to evaluate them
                if moves - self.state_0.g < self.lookahead_depth:
                    for action in state.actions():
                        self.state_graph.add_node((state, action))
                        self.state_graph.add_edge(state, (state, action))                
                        for (next_state, next_state_prob) in self.generate_next_states(state, action):
                            self.add_state(state, v, openlist, moves, action, next_state, next_state_prob)
                else: # Lookahead limit reached - need to use heuristic
                    v[state]=moves+self.h(state)
                    policy[state]=None
                    
    def iterate_state(self, state,v):
        if self.is_goal(state):
            return (self.g,None)
        if self.state_graph.order(state)==0: # Tip node (no actions or lookahead limit) - use heuristic
            return (state.g+self.h(state),None)          
        else:
            return cliques.mdp.heuristicMdp.heuristicMdp.iterate_state(self,state,v)

